from dataset import MakeDataSet
import numpy as np
import os
from PIL import Image

def img_mean_and_std():
    image_channels = 3
    folder_path = MakeDataSet.source_path
    # 初始化积累变量
    total_pixels = 0
    # 如果是RGB图像，需要3个通道的均值和方差
    sum_normalized_pixel_values = np.zeros(image_channels)
    # 遍历所有的图片文件
    for root,dirs,files in os.walk(folder_path):
        for filename in files:
            if filename.endswith(MakeDataSet.image_suffix_list):
                try:
                    image_path = os.path.join(root,filename)
                    image = Image.open(image_path)
                    image_array = np.array(image)
                    ## 归一化像素值到0-1之间
                    normalized_image_array = image_array/255.0
                    total_pixels+=normalized_image_array.size
                    sum_normalized_pixel_values+=np.sum(normalized_image_array,axis=(0,1))
                except Exception:
                    pass

    ## 计算均值和方差
    mean = sum_normalized_pixel_values / total_pixels

    sum_squared_diff = np.zeros(image_channels)
    for root,dirs,files in os.walk(folder_path):
        for filename in files:
            if filename.endswith(MakeDataSet.image_suffix_list):
                try:
                    image_path = os.path.join(root, filename)
                    image = Image.open(image_path)
                    image_array = np.array(image)
                    # 归一化像素在0-1之间
                    normalized_image_array = image_array / 255.0
                    print(normalized_image_array.shape)
                    print(mean.shape)
                    print(image_path)
                    diff = (normalized_image_array-mean)**2
                    sum_squared_diff += np.sum(diff,axis=(0,1))
                except Exception:
                    print(f'{image_path}图片异常')

    variance = sum_squared_diff/total_pixels
    print(f'均值={mean},均方差={variance}')
    ## 返回均值和均值方差
    return mean,variance

if __name__ == '__main__':
    img_mean_and_std()